import pandas as pd
from sklearn.model_selection import train_test_split
import os
import json
from datetime import datetime


def convert_dataset_to_tsv(
        input_csv,
        output_dir,
        input_text_column="review",  # 输入文件的文本列名
        input_label_column="label",  # 输入文件的标签列名
        output_text_column="text_a",  # 输出文件的文本列名
        output_label_column="label",  # 输出文件的标签列名
        test_size=0.1,
        dev_size=0.1,
        random_seed=42
):
    """
    数据集格式转换工具（CSV -> KBert TSV）
    功能：
    1. 自动将(review, label)格式转为(text_a, label)格式
    2. 保留完整元数据
    3. 处理特殊符号
    """
    os.makedirs(output_dir, exist_ok=True)

    try:
        # 读取CSV（自动处理编码）
        df = pd.read_csv(input_csv, encoding='utf-8')
        print(f"✅ 成功读取输入文件: {os.path.basename(input_csv)}")
        print(f"  样本数: {len(df):,} | 原始列名: {list(df.columns)}")

        # 验证输入列是否存在
        for col in [input_text_column, input_label_column]:
            if col not in df.columns:
                raise ValueError(f"❌ 输入文件中缺少必要列: '{col}'")

        # 重命名列（review -> text_a）
        df = df.rename(columns={
            input_text_column: output_text_column,
            input_label_column: output_label_column
        })

        # 划分数据集
        train_df, test_df = train_test_split(
            df,
            test_size=test_size,
            random_state=random_seed,
            stratify=df[output_label_column]
        )
        train_df, dev_df = train_test_split(
            train_df,
            test_size=dev_size / (1 - test_size),
            random_state=random_seed,
            stratify=train_df[output_label_column]
        )

        # 保存TSV（带目标表头）
        for split_name, split_data in [
            ('train', train_df),
            ('dev', dev_df),
            ('test', test_df)
        ]:
            output_path = os.path.join(output_dir, f"{split_name}.tsv")
            split_data[[output_label_column, output_text_column]].to_csv(
                output_path,
                sep='\t',
                index=False,
                header=True  # 保留表头
            )
            print(f"  已生成: {output_path} ({len(split_data):,}条)")

        print("\n✨ 转换完成！输出结构：")
        print(f"{output_dir}/\n├── train.tsv\n├── dev.tsv\n└── test.tsv")
        print(f"\n类别分布:\n{df[output_label_column].value_counts().to_string()}")

    except Exception as e:
        print(f"❌ 处理失败: {str(e)}")
        if 'df' in locals():
            print("数据样例：")
            print(df.head(2))


if __name__ == "__main__":
    # 配置参数
    config = {
        "input_csv": "E:\词向量\dataset\中文二分类\ChnSentiCorp_htl_all（中文二分类）.csv",  # 输入文件路径
        "output_dir": "./crop_tsv",  # 输出目录
        "input_text_column": "review",  # 输入文件的文本列名
        "input_label_column": "label",  # 输入文件的标签列名
        "test_size": 0.1,  # 测试集比例
        "dev_size": 0.1,  # 验证集比例
        "random_seed": 42  # 随机种子
    }

    convert_dataset_to_tsv(**config)